

const tf = require("@tensorflow/tfjs");
const csv = require('node-csv').createParser();
const util = require('../util.js');

let outputSize = 3	// 预测值onehot长度
let inputShape = [4]	//data长度
let lr = 0.01	
let batchSize = 32	//每次训练多少个
let total_episode = 51	//总训练次数
let test_num = 30	//多少个数据作为 测试数据
let csvPath = './datasets/Iris.csv' //csv数据文件位置
/*
*	定义model
*/
const model = tf.sequential({
	     layers: [
	     	tf.layers.dense({
	     		units: outputSize * 3, 	//output shape
	     		inputShape: inputShape, 	//input shape
	     		kernelInitializer: 'randomNormal', 	//kernel的初始化方式
	     		biasInitializer: 'ones', 	//bias初始化方式
	     		activation: 'relu'	//激励函数
	     	}),
	     	tf.layers.dropout({rate: 0.2}),	//droput层
	     	tf.layers.dense({units: outputSize, inputShape: [outputSize * 3], kernelInitializer: 'randomNormal', biasInitializer: 'ones'})
	     ]
	});
model.compile({
	optimizer: tf.train.adam(lr), //优化器 
	loss: 'meanSquaredError'	//loss计算方式
});

/*
*	转换label为onehot格式
*/
const transformLabel = (labels, names, value) => {
	const tampleLabel = [0, 0, 0]
	tampleLabel[names[value.pop()]] = 1
	labels.push(tampleLabel)
}

/*
*	处理csv读取出来的数据，返回[dict，labels，data]
*/
const parseData = data => {
	var count = 0; const names = {}; const labels = []; 
	data.shift()
	data.map(value => {
		const name = value[value.length - 1]
		names[name] = names[name] || names[name] == 0 ? names[name] : count++
		transformLabel(labels, names, value)
		value.splice(0, 1)
	})
	return {names: names, labels: labels, data: data}
}


/*
*	训练流程
*/
const train = async parsedData => {
	const names = parsedData['names'], labels = parsedData['labels'], data = parsedData['data']
	const sample_input = util.sample(data, labels, test_num)
	const test_data = tf.tensor(sample_input['test_data']), test_labels = tf.tensor(sample_input['test_labels'])
	const train_data = tf.tensor(sample_input['train_data']), train_labels = tf.tensor(sample_input['train_labels'])
	var h;
	for (let i = 1; i <= total_episode ; ++i) {
	   	h = await model.fit(train_data, train_labels, {batchSize: batchSize, epochs: 1});
	   	if (i % 10 == 1) {
	   		console.log("Loss after Epoch " + i + " : " + h.history.loss[0]);
	   	}
	}
	return {test_data: test_data, test_labels: test_labels}
}

/*
*	测试流程
*/
const tester = _test => {
	var test_data = _test['test_data'], test_labels = _test['test_labels']
	var predict = model.predict(test_data)
	predict = predict.argMax(1).dataSync()
	test_labels = test_labels.argMax(1).dataSync()
	console.log('predict labels :\t' + predict)
	console.log('true labels :\t\t' + test_labels)
	var conut_right = 0;
	var total_count = 0;
	predict.forEach((pred, i) => {
		total_count ++
		conut_right += pred == test_labels[i] ? 1 : 0
	})
	console.log('total_count :\t\t' + total_count)
	console.log('conut_right :\t\t' + conut_right)
	console.log('Correct rate :\t\t' + ((conut_right / total_count) * 100) + '%')
}

/*
*	读取文件并开始训练以及测试
*/
csv.parseFile(csvPath, async (err, data) => {
	var test_datas = await train(parseData(data))
	tester(test_datas)
});
